import asyncio
import re
from douyin.douyin_pb2 import PushFrame, Response, ChatMessage
import gzip
import httpx
from models import User, live_data
import random
from settings import websocket_connections
from ai_tools import ai
import time
from functools import partial
import subprocess
# 锁定输出utf-8
subprocess.Popen = partial(subprocess.Popen, encoding='utf-8')
import execjs
import hashlib
import websockets


# 对每一个ws连接发送数据处理结果的消息
async def broadcast_message(secuid: str, message: str):
    for websocket in websocket_connections:
        try:
            await websocket.send_json({"secuid": secuid, "msg": message})
        except Exception as e:
            print(f"Error sending message to client: {e}")
            websocket_connections.remove(websocket)


# 触发关键词的消息进行处理
async def data_process(message: ChatMessage,user_id: str , room_id: str, word: str, rewords: list):
    is_active_ai = (await User.get(uid=user_id).values("is_active_ai"))['is_active_ai']
    # 判断是否启用了ai
    if is_active_ai:
        result, ai_evaluation = await ai(message.content)
    else:
        result = False
        ai_evaluation = ""
    # 不管调用不调用ai都存数据库
    await live_data.create(room_id=room_id, nickName=message.user.nickName,
                           sec_uid=message.user.sec_uid, keyword=word, message=message.content,
                           ai_evaluation=ai_evaluation,uid = user_id)
    # 基于ws给客户端发送消息
    if result:
        await broadcast_message(message.user.sec_uid, random.choice(rewords))

# 获取signature
def get_signature(room_id,user_unique_id):
    f = open('douyin\getsig.js', mode='r', encoding='utf-8')
    js = f.read()
    f.close()
    JS = execjs.compile(js)
    hash_object = hashlib.md5()
    a = (f'live_id=1,aid=6383,version_code=180800,webcast_sdk_version=1.0.14-beta.0,room_id={room_id},sub_room_id=,'
    f'sub_channel_id=,did_rule=3,user_unique_id={user_unique_id},device_platform=web,device_type=,ac=,identity=audience')
    # 将输入字符串编码为字节，MD5需要字节类型的数据
    byte_string = a.encode('utf-8')

    # 更新hash对象，使其处理输入的字节数据
    hash_object.update(byte_string)

    # 获取十六进制格式的哈希值
    hex_dig = hash_object.hexdigest()
    sign = JS.call("get_signature", hex_dig)
    return sign

# 获取直播间参数
async def fetch_live_room_info(url):
    cookies = {
        'passport_csrf_token': '5adc63a3f5c316b94d27e46fa63c76a3',
        'passport_csrf_token_default': '5adc63a3f5c316b94d27e46fa63c76a3',
        'bd_ticket_guard_client_web_domain': '2',
        'SEARCH_RESULT_LIST_TYPE': '%22single%22',
        'UIFID_TEMP': '3c3e9d4a635845249e00419877a3730e2149197a63ddb1d8525033ea2b3354c2ec55162f7fdf06585b70ef346752ab878a13b38a3dfe72cbbc0ac1f47838a5341120a9cfc199436ddd6a28f15c6c3d88cf1990736db1afa6c62e8393dd4a89dfa232d67e27afec8e3d115fd1076a1602',
        'live_use_vvc': '%22false%22',
        'xgplayer_user_id': '834977490424',
        'fpk1': 'U2FsdGVkX1/BF1RdWdjXkAGf8ZbF5uYkSqQm2HyWMVqA7Cd9HsZBnCXEmW8W4E2KhgaTd901aXAncgTHFAblBw==',
        'fpk2': 'f1f6b29a6cc1f79a0fea05b885aa33d0',
        'UIFID': '3c3e9d4a635845249e00419877a3730e2149197a63ddb1d8525033ea2b3354c2ab4b6950de9ed210ef7a935110b4430f086100d4e5a3a63a6d568a10534dfa920db4eb240b501948453d79039fc2068668bd9b93b9c589ba963620debb478a0225da18776baa9da7664cb8ffaa52f8e34cf68cbae5ecbfb3e7006801498715aa2ecba5dbc99a94c99be2733965ef9cc69d42df9f322a4519463b05513c5b50b3',
        'n_mh': 'Gh_pdaCQRNbi3LgWLamH4PjjVQ3_R5SrKjnmCalg2vc',
        '_bd_ticket_crypt_doamin': '2',
        '_bd_ticket_crypt_cookie': 'f896dac0fb946d4b718d8ba1b81670a7',
        '__security_server_data_status': '1',
        'store-region': 'cn-jx',
        'store-region-src': 'uid',
        'ttwid': '1%7CbgYRNuE7he9oJgE-ZLFK6XytpYBRWyKwfTvU1AEhg8s%7C1719928171%7C8fa43855dc56396085f804b10f28fb58575954ff65b7575dbd1c5a7cf10a21e5',
        'odin_tt': '162e3f7243e0faa2f7242f5219e9e97c3ea5e721d3dee41710a4b35c064f39099a47a2144a8ef81bf9ec68742f55127e570828f42ee19aebda54c366021e99708456e509ab2490d1d27b919b24736c1b',
        'bd_ticket_guard_client_data': 'eyJiZC10aWNrZXQtZ3VhcmQtdmVyc2lvbiI6MiwiYmQtdGlja2V0LWd1YXJkLWl0ZXJhdGlvbi12ZXJzaW9uIjoxLCJiZC10aWNrZXQtZ3VhcmQtcmVlLXB1YmxpYy1rZXkiOiJCRk5aZUs0UnJJbVJsUWxQUG16c1dxZ1VxWTFRbHZEbEFNS21FTGNpR2JGK25ZY1FkWEFsLzM3OWhiZ1E5VFFVdGczYWIyRnNWOXQ1TnE3KzRPck5jd289IiwiYmQtdGlja2V0LWd1YXJkLXdlYi12ZXJzaW9uIjoxfQ%3D%3D',
        'FORCE_LOGIN': '%7B%22videoConsumedRemainSeconds%22%3A180%7D',
        'volume_info': '%7B%22isMute%22%3Atrue%2C%22isUserMute%22%3Atrue%7D',
        '__live_version__': '%221.1.2.2525%22',
        'has_avx2': 'null',
        'device_web_cpu_core': '8',
        'device_web_memory_size': '8',
        'csrf_session_id': '66c5397e51888e9974829dd5fd7faebf',
        'download_guide': '%223%2F20240803%2F0%22',
        'pwa2': '%220%7C0%7C2%7C0%22',
        'xg_device_score': '7.628058823529412',
        'live_can_add_dy_2_desktop': '%221%22',
        'IsDouyinActive': 'false',
        '__ac_nonce': '066ae18a700287e94166d',
        '__ac_signature': '_02B4Z6wo00f01pu6jJAAAIDDR8MTJAY2txKbmogAAMBM64',
        '__ac_referer': url,
    }

    headers = {
        'accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        'accept-language': 'zh-CN,zh;q=0.9',
        'cache-control': 'max-age=0',
        # 'cookie': 'passport_csrf_token=5adc63a3f5c316b94d27e46fa63c76a3; passport_csrf_token_default=5adc63a3f5c316b94d27e46fa63c76a3; bd_ticket_guard_client_web_domain=2; SEARCH_RESULT_LIST_TYPE=%22single%22; UIFID_TEMP=3c3e9d4a635845249e00419877a3730e2149197a63ddb1d8525033ea2b3354c2ec55162f7fdf06585b70ef346752ab878a13b38a3dfe72cbbc0ac1f47838a5341120a9cfc199436ddd6a28f15c6c3d88cf1990736db1afa6c62e8393dd4a89dfa232d67e27afec8e3d115fd1076a1602; live_use_vvc=%22false%22; xgplayer_user_id=834977490424; fpk1=U2FsdGVkX1/BF1RdWdjXkAGf8ZbF5uYkSqQm2HyWMVqA7Cd9HsZBnCXEmW8W4E2KhgaTd901aXAncgTHFAblBw==; fpk2=f1f6b29a6cc1f79a0fea05b885aa33d0; UIFID=3c3e9d4a635845249e00419877a3730e2149197a63ddb1d8525033ea2b3354c2ab4b6950de9ed210ef7a935110b4430f086100d4e5a3a63a6d568a10534dfa920db4eb240b501948453d79039fc2068668bd9b93b9c589ba963620debb478a0225da18776baa9da7664cb8ffaa52f8e34cf68cbae5ecbfb3e7006801498715aa2ecba5dbc99a94c99be2733965ef9cc69d42df9f322a4519463b05513c5b50b3; n_mh=Gh_pdaCQRNbi3LgWLamH4PjjVQ3_R5SrKjnmCalg2vc; _bd_ticket_crypt_doamin=2; _bd_ticket_crypt_cookie=f896dac0fb946d4b718d8ba1b81670a7; __security_server_data_status=1; store-region=cn-jx; store-region-src=uid; ttwid=1%7CbgYRNuE7he9oJgE-ZLFK6XytpYBRWyKwfTvU1AEhg8s%7C1719928171%7C8fa43855dc56396085f804b10f28fb58575954ff65b7575dbd1c5a7cf10a21e5; odin_tt=162e3f7243e0faa2f7242f5219e9e97c3ea5e721d3dee41710a4b35c064f39099a47a2144a8ef81bf9ec68742f55127e570828f42ee19aebda54c366021e99708456e509ab2490d1d27b919b24736c1b; bd_ticket_guard_client_data=eyJiZC10aWNrZXQtZ3VhcmQtdmVyc2lvbiI6MiwiYmQtdGlja2V0LWd1YXJkLWl0ZXJhdGlvbi12ZXJzaW9uIjoxLCJiZC10aWNrZXQtZ3VhcmQtcmVlLXB1YmxpYy1rZXkiOiJCRk5aZUs0UnJJbVJsUWxQUG16c1dxZ1VxWTFRbHZEbEFNS21FTGNpR2JGK25ZY1FkWEFsLzM3OWhiZ1E5VFFVdGczYWIyRnNWOXQ1TnE3KzRPck5jd289IiwiYmQtdGlja2V0LWd1YXJkLXdlYi12ZXJzaW9uIjoxfQ%3D%3D; FORCE_LOGIN=%7B%22videoConsumedRemainSeconds%22%3A180%7D; volume_info=%7B%22isMute%22%3Atrue%2C%22isUserMute%22%3Atrue%7D; __live_version__=%221.1.2.2525%22; has_avx2=null; device_web_cpu_core=8; device_web_memory_size=8; csrf_session_id=66c5397e51888e9974829dd5fd7faebf; download_guide=%223%2F20240803%2F0%22; pwa2=%220%7C0%7C2%7C0%22; xg_device_score=7.628058823529412; live_can_add_dy_2_desktop=%221%22; IsDouyinActive=false; __ac_nonce=066ae18a700287e94166d; __ac_signature=_02B4Z6wo00f01pu6jJAAAIDDR8MTJAY2txKbmogAAMBM64; __ac_referer=https://live.douyin.com/80017709309',
        'priority': 'u=0, i',
        'referer': url,
        'sec-ch-ua': '"Not)A;Brand";v="99", "Google Chrome";v="127", "Chromium";v="127"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Windows"',
        'sec-fetch-dest': 'document',
        'sec-fetch-mode': 'navigate',
        'sec-fetch-site': 'same-origin',
        'upgrade-insecure-requests': '1',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36',
    }
    async with httpx.AsyncClient(verify=False) as client:
        response = await client.get(
            url = url,
            headers=headers,
            cookies=cookies
        )
        html = response.text
    # 提取所需的字段
    room_id = re.findall(r'roomId\\":\\"(\d+)\\"', html)[0]
    user_unique_id = re.findall(r'user_unique_id\\":\\"(\d+)\\"', html)[0]

    first_time = int(time.time() * 1000)
    sign = get_signature(room_id,user_unique_id)
    fetch_time = first_time + random.randint(80,120)

    wss_url = (f"wss://webcast5-ws-web-hl.douyin.com/webcast/im/push/v2/?app_name=douyin_web&version_code=180800"
               f"&webcast_sdk_version=1.0.14-beta.0&update_version_code=1.0.14-beta.0&compress=gzip&device_platform"
               f"=web&cookie_enabled=true&screen_width=1920&screen_height=1080&browser_language=zh-CN"
               f"&browser_platform=Win32&browser_name=Mozilla&browser_version=5.0%20("
               f"Windows%20NT%2010.0;%20Win64;%20x64)%20AppleWebKit/537.36%20(KHTML,"
               f"%20like%20Gecko)%20Chrome/126.0.0.0%20Safari/537.36&browser_online=true&tz_name=Asia/Shanghai&cursor"
               f"=t-1721376268998_r-1_d-1_u-1_fh-7393254521105667110&internal_ext=internal_src:dim|wss_push_room_id"
               f":{room_id}|wss_push_did:7393234104294147634|first_req_ms:{first_time}|fetch_time"
               f":{fetch_time}|seq:1|wss_info:0-{fetch_time}-0-0|wrds_v:7393254775170532758&host=https://live"
               f".douyin.com&aid=6383&live_id=1&did_rule=3&endpoint=live_pc&support_wrds=1&user_unique_id"
               f"={user_unique_id}&im_path=/webcast/im/fetch/&identity=audience&need_persist_msg_count=15"
               f"&insert_task_id=&live_reason=&room_id={room_id}&heartbeatDuration=0&signature={sign}")

    ttwid = response.cookies.get('ttwid')
    return {"url": wss_url, "ttwid": ttwid}



class douyin:

    def __init__(self,user_id:str) -> None:
        self.user_id = user_id
        self.tasks: dict[str, asyncio.Task[None]] = {}
        self.keywords: dict[str, list] = {}
        self.rewords: dict[str, list] = {}

    async def on_message(self, ws, content, room_id):
        frame = PushFrame()
        frame.ParseFromString(content)

        # 对PushFrame的 payload 内容进行gzip解压
        origin_bytes = gzip.decompress(frame.payload)

        # 根据Response+gzip解压数据，生成数据对象
        response = Response()
        response.ParseFromString(origin_bytes)

        if response.needAck:
            s = PushFrame()
            s.payloadType = "ack"
            s.payload = response.internalExt.encode('utf-8')
            s.logId = frame.logId
        
            await ws.send(s.SerializeToString())

        # 获取数据内容（需根据不同method，使用不同的结构对象对 数据 进行解析）
        #   注意：此处只处理 WebcastChatMessage ，其他处理方式都是类似的。
        for item in response.messagesList:
            if item.method != "WebcastChatMessage":
                continue
            message = ChatMessage()
            message.ParseFromString(item.payload)
            # 对数据进行筛选分析存入数据库
            for word in self.keywords[room_id]:
                if word in message.content:
                    data = await live_data.filter(room_id=room_id, sec_uid=message.user.sec_uid).values('message')
                    msg_list = [msg['message'] for msg in data]
                    if message.content not in msg_list:
                        # 下面进行进一步的评论分析操作
                        await asyncio.create_task(data_process(message, self.user_id, room_id, word, self.rewords[room_id]))
                        break

    async def receive_messages(self, url, room_id,tw):
        print("开始一个任务")
        # page = await self.context.new_page()
        # self.task_page[room_id] = page

        # async def handle_websocket(content):
        #     await self.on_message(content, room_id)

        # page.on('websocket', lambda ws: ws.on("framereceived", handle_websocket))
        # await page.goto(url)

        # await page.wait_for_selector("//div[text() = '在线观众']")
        # while True:
        #     try:
        #         await asyncio.sleep(random.randint(15, 30))
        #         await page.click("//div[text() = '在线观众']")
        #     except Exception as e:
        #         print(e)
        #         continue

        headers = {
                "connection":"Upgrade",
                "host":"webcast5-ws-web-hl.douyin.com",
                "origin":"https://live.douyin.com",
                "sec-websocket-version":"13",
                "upgrade":"websocket",
                "user-agent":"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36",
                "Cookie":f"ttwid={tw}"
            }

        async with websockets.connect(url, extra_headers=headers) as websocket:
            try:
                while 1:
                    message = await websocket.recv()
                    # 在这里处理接收到的message
                    await self.on_message(websocket, message, room_id)
                    await asyncio.sleep(0.5)
            except websockets.exceptions.ConnectionClosedError:
                print("Connection closed unexpectedly, retrying...")
            except Exception as e:
                print(e)
 

    # 添加一个task任务
    async def add_task(self, url: str, keyword_list: list[str], reword_list: list[str]):
        room_id = url.split('/')[-1]
        data = await fetch_live_room_info(url=url)
        # 不在任务列表里就创建任务
        if room_id not in self.tasks.keys():
            new_task = asyncio.create_task(self.receive_messages(url=data['url'], room_id=room_id,tw=data['ttwid']))
            self.tasks[room_id] = new_task
            self.keywords[room_id] = keyword_list
            self.rewords[room_id] = reword_list

    # 删除指定的task任务
    async def del_task(self, room_id: str):
        self.tasks[room_id].cancel()
        del self.tasks[room_id]
        del self.keywords[room_id]
        del self.rewords[room_id]


class Task_user:
    def __init__(self):
        self.User: dict[str, douyin] = {}

    def query_task(self, user_id):
        if user_id not in self.User:
            return {}
        else:
            return self.User[user_id].tasks

    async def add_task(self, user_id, url: str, keyword_list: list[str], reword_list: list[str]):
        if user_id not in self.User:
            self.User[user_id] = douyin(user_id=user_id)
        await self.User[user_id].add_task(url, keyword_list, reword_list)

    async def del_task(self, user_id: str,room_id: str):
        await self.User[user_id].del_task(room_id)







if __name__ == "__main__":
    ob_dy = douyin()
    asyncio.run(ob_dy.add_task("https://live.douyin.com/265366259141",['1','2'],['1','2']))